package com.demo.newLSTM;

import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.GradientNormalization;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.*;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.learning.config.Adam;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.util.ArrayList;
import java.util.List;

/**
 * Java编写的LSTM短期负荷预测 加入 季节、节假日、负荷特性、气象因素、政策因素、经济因素、负荷管理措施等影响因素
 */
public class ShortTermLoadForecasting1 {
    public static void main(String[] args) {
        // 准备训练数据
        List<DataSet> trainData = prepareTrainingData();

        // 创建LSTM网络模型
        int numInputs = 7; // 输入特征数量
        int numOutputs = 1; // 输出负荷数量
        int lstmLayerSize = 128; // LSTM层大小
        int numHiddenLayers = 2; // 隐藏层数量

        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
                .seed(123)
                .optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT)
                .weightInit(WeightInit.XAVIER)
                .updater(new Adam(0.01))
                .gradientNormalization(GradientNormalization.ClipElementWiseAbsoluteValue)
                .gradientNormalizationThreshold(0.5)
                .list()
                .layer(0, new LSTM.Builder()
                        .nIn(numInputs)
                        .nOut(lstmLayerSize)
                        .activation(Activation.TANH)
                        .build())
                .layer(1, new RnnOutputLayer.Builder(LossFunctions.LossFunction.MSE)
                        .activation(Activation.IDENTITY)
                        .nIn(lstmLayerSize)
                        .nOut(numOutputs)
                        .build())
                .build();

        MultiLayerNetwork model = new MultiLayerNetwork(conf);
        model.init();

        // 训练模型
        int numEpochs = 100;
        for (int i = 0; i < numEpochs; i++) {
            for (DataSet ds : trainData) {
                model.fit(ds);
            }
        }

        // 使用模型进行预测
        INDArray input = Nd4j.create(new double[]{0.8, 0.5, 0.6, 0.7, 0.9, 0.8, 0.6}).reshape(1, numInputs);
        INDArray output = model.output(input);

        System.out.println("预测负荷值：" + output);
    }

    private static List<DataSet> prepareTrainingData() {
        // 这里假设有一些训练数据，每个样本包含7个输入特征和1个输出负荷值
        List<DataSet> trainData = new ArrayList<>();
        trainData.add(new DataSet(Nd4j.create(new double[]{0.7, 0.8, 0.9, 0.6, 0.5, 0.7, 0.6}), Nd4j.create(new double[]{0.85})));
        trainData.add(new DataSet(Nd4j.create(new double[]{0.6, 0.5, 0.8, 0.7, 0.9, 0.6, 0.7}), Nd4j.create(new double[]{0.75})));
        trainData.add(new DataSet(Nd4j.create(new double[]{0.8, 0.6, 0.7, 0.9, 0.5, 0.7, 0.8}), Nd4j.create(new double[]{0.95})));

        return trainData;
    }
}
